/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

package vip.zhenzicheng.rocketmq.api.batch;

import org.apache.rocketmq.client.producer.DefaultMQProducer;
import org.apache.rocketmq.common.message.Message;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

public class SplitBatchProducer {

  public static void main(String[] args) throws Exception {

    DefaultMQProducer producer = new DefaultMQProducer("BatchProducerGroupName");
    producer.start();

    //large batch
    String topic = "BatchTest";
    List<Message> messages = new ArrayList<>(100 * 1000);
    for (int i = 0; i < 100 * 1000; i++) {
      messages.add(new Message(topic, "Tag", "OrderID" + i, ("Hello world " + i).getBytes()));
    }
//        producer.send(messages);

    //split the large batch into small ones:
    ListSplitter splitter = new ListSplitter(messages);
    while (splitter.hasNext()) {
      List<Message> listItem = splitter.next();
      producer.send(listItem);
    }
    producer.shutdown();
  }

}

class ListSplitter implements Iterator<List<Message>> {
  private final List<Message> messages;
  private final int sizeLimit = 1000 * 1000;
  private int currIndex;

  public ListSplitter(List<Message> messages) {
    this.messages = messages;
  }

  @Override
  public boolean hasNext() {
    return currIndex < messages.size();
  }

  @Override
  public List<Message> next() {
    int nextIndex = currIndex;
    int totalSize = 0;
    for (; nextIndex < messages.size(); nextIndex++) {
      Message message = messages.get(nextIndex);
      int tmpSize = message.getTopic().length() + message.getBody().length;
      Map<String, String> properties = message.getProperties();
      for (Map.Entry<String, String> entry : properties.entrySet()) {
        tmpSize += entry.getKey().length() + entry.getValue().length();
      }
      tmpSize = tmpSize + 20; //for log overhead
      if (tmpSize > sizeLimit) {
        //it is unexpected that single message exceeds the sizeLimit
        //here just let it go, otherwise it will block the splitting process
        if (nextIndex - currIndex == 0) {
          //if the next sublist has no element, add this one and then break, otherwise just break
          nextIndex++;
        }
        break;
      }
      if (tmpSize + totalSize > sizeLimit) {
        break;
      } else {
        totalSize += tmpSize;
      }

    }
    List<Message> subList = messages.subList(currIndex, nextIndex);
    currIndex = nextIndex;
    return subList;
  }

  @Override
  public void remove() {
    throw new UnsupportedOperationException("Not allowed to remove");
  }
}
